# -*- coding: utf-8 -*-
"""
代理管理模块
支持隧道代理和IP代理池管理
"""

import requests
import random
import time
import logging
from typing import List, Dict, Optional, Tuple
from config.base_settings import PROXY_CONFIG, REDIS_CONFIG, REDIS_DBS, REDIS_KEYS

logger = logging.getLogger(__name__)


class ProxyManager:
    """代理管理器"""
    
    def __init__(self):
        """初始化代理管理器"""
        self.proxy_type = PROXY_CONFIG['PROXY_TYPE']
        self.tunnel_proxy = PROXY_CONFIG.get('TUNNEL_PROXY', {})
        self.proxy_pool = PROXY_CONFIG.get('PROXY_POOL', [])
        self.current_proxy_index = 0
        self.failed_proxies = set()  # 记录失败的代理
        
        # 初始化代理池到Redis
        if self.proxy_type == 'pool' and self.proxy_pool:
            self._init_proxy_pool()
    
    def _init_proxy_pool(self):
        """初始化代理池到Redis"""
        try:
            import redis
            redis_client = redis.Redis(**REDIS_CONFIG, db=REDIS_DBS['proxy_pool'])
            
            # 清空现有代理池
            redis_client.delete(REDIS_KEYS['PROXY_POOL'])
            
            # 添加代理到池中
            for proxy in self.proxy_pool:
                redis_client.sadd(REDIS_KEYS['PROXY_POOL'], proxy)
            
            logger.info(f"代理池初始化完成，共 {len(self.proxy_pool)} 个代理")
            
        except Exception as e:
            logger.error(f"初始化代理池失败: {e}")
    
    def get_proxy(self) -> Optional[Dict[str, str]]:
        """获取代理配置"""
        if not PROXY_CONFIG['USE_PROXY']:
            return None
        
        if self.proxy_type == 'tunnel':
            return self.tunnel_proxy
        elif self.proxy_type == 'pool':
            return self._get_pool_proxy()
        
        return None
    
    def _get_pool_proxy(self) -> Optional[Dict[str, str]]:
        """从代理池中获取代理"""
        try:
            import redis
            redis_client = redis.Redis(**REDIS_CONFIG, db=REDIS_DBS['proxy_pool'])
            
            # 获取可用代理
            available_proxies = redis_client.smembers(REDIS_KEYS['PROXY_POOL'])
            available_proxies = [p.decode() if isinstance(p, bytes) else p for p in available_proxies]
            
            if not available_proxies:
                logger.warning("代理池为空")
                return None
            
            # 过滤掉失败的代理
            working_proxies = [p for p in available_proxies if p not in self.failed_proxies]
            
            if not working_proxies:
                logger.warning("所有代理都已失败，重置失败列表")
                self.failed_proxies.clear()
                working_proxies = available_proxies
            
            # 选择代理
            if PROXY_CONFIG.get('PROXY_ROTATION', True):
                proxy = working_proxies[self.current_proxy_index % len(working_proxies)]
                self.current_proxy_index += 1
            else:
                proxy = random.choice(working_proxies)
            
            return {
                'http': proxy,
                'https': proxy
            }
            
        except Exception as e:
            logger.error(f"获取代理池代理失败: {e}")
            return None
    
    def test_proxy(self, proxy: Dict[str, str], timeout: int = 10) -> bool:
        """测试代理是否可用"""
        try:
            test_url = "https://www.baidu.com"
            response = requests.get(
                test_url, 
                proxies=proxy, 
                timeout=timeout,
                headers={'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'}
            )
            
            if response.status_code == 200:
                logger.info(f"代理测试成功: {proxy}")
                return True
            else:
                logger.warning(f"代理测试失败，状态码: {response.status_code}")
                return False
                
        except Exception as e:
            logger.warning(f"代理测试失败: {e}")
            return False
    
    def mark_proxy_failed(self, proxy: Dict[str, str]):
        """标记代理为失败"""
        if self.proxy_type == 'pool':
            proxy_url = proxy.get('http') or proxy.get('https')
            if proxy_url:
                self.failed_proxies.add(proxy_url)
                logger.warning(f"标记代理失败: {proxy_url}")
    
    def get_working_proxy(self) -> Optional[Dict[str, str]]:
        """获取可用的代理"""
        max_retries = PROXY_CONFIG.get('PROXY_RETRY', 3)
        return self.get_proxy()
        
        for attempt in range(max_retries):
            proxy = self.get_proxy()
            if not proxy:
                return None
            
            if self.test_proxy(proxy, PROXY_CONFIG.get('PROXY_TIMEOUT', 10)):
                return proxy
            else:
                self.mark_proxy_failed(proxy)
                time.sleep(1)  # 等待1秒后重试
        
        logger.error("无法获取可用代理")
        return None
    
    def is_proxy_enabled(self) -> bool:
        """检查是否启用代理"""
        return PROXY_CONFIG['USE_PROXY']
    
    def get_proxy_stats(self) -> Dict:
        """获取代理统计信息"""
        try:
            import redis
            redis_client = redis.Redis(**REDIS_CONFIG, db=REDIS_DBS['proxy_pool'])
            
            total_proxies = redis_client.scard(REDIS_KEYS['PROXY_POOL'])
            failed_count = len(self.failed_proxies)
            working_count = total_proxies - failed_count
            
            return {
                'total_proxies': total_proxies,
                'working_proxies': working_count,
                'failed_proxies': failed_count,
                'proxy_type': self.proxy_type
            }
            
        except Exception as e:
            logger.error(f"获取代理统计失败: {e}")
            return {}
    
    def reset_failed_proxies(self):
        """重置失败代理列表"""
        self.failed_proxies.clear()
        logger.info("已重置失败代理列表")
    
    def add_proxy_to_pool(self, proxy_url: str):
        """添加代理到池中"""
        if self.proxy_type == 'pool':
            try:
                import redis
                redis_client = redis.Redis(**REDIS_CONFIG, db=REDIS_DBS['proxy_pool'])
                redis_client.sadd(REDIS_KEYS['PROXY_POOL'], proxy_url)
                logger.info(f"添加代理到池中: {proxy_url}")
            except Exception as e:
                logger.error(f"添加代理失败: {e}")
    
    def remove_proxy_from_pool(self, proxy_url: str):
        """从池中移除代理"""
        if self.proxy_type == 'pool':
            try:
                import redis
                redis_client = redis.Redis(**REDIS_CONFIG, db=REDIS_DBS['proxy_pool'])
                redis_client.srem(REDIS_KEYS['PROXY_POOL'], proxy_url)
                logger.info(f"从池中移除代理: {proxy_url}")
            except Exception as e:
                logger.error(f"移除代理失败: {e}")


class RequestHelper:
    """请求辅助类"""
    
    def __init__(self, proxy_manager: ProxyManager):
        """初始化请求辅助类"""
        self.proxy_manager = proxy_manager
        self.session = None
    
    def make_request(self, url: str, params: Dict = None, headers: Dict = None, 
                    timeout: int = 30, retries: int = 3) -> Optional[requests.Response]:
        """发起HTTP请求"""
        for attempt in range(retries):
            try:
                # 获取代理
                proxy = self.proxy_manager.get_working_proxy() if self.proxy_manager.is_proxy_enabled() else None
                
                # 发起请求
                response = requests.get(
                    url,
                    params=params,
                    headers=headers,
                    proxies=proxy,
                    timeout=timeout
                )
                
                response.raise_for_status()
                return response
                
            except requests.exceptions.ProxyError as e:
                logger.warning(f"代理错误 (尝试 {attempt + 1}/{retries}): {e}")
                if proxy:
                    self.proxy_manager.mark_proxy_failed(proxy)
                time.sleep(2 ** attempt)
                
            except requests.exceptions.RequestException as e:
                logger.warning(f"请求失败 (尝试 {attempt + 1}/{retries}): {e}")
                time.sleep(2 ** attempt)
        
        logger.error(f"请求最终失败: {url}")
        return None